package com.feidee.fd.sml.algorithm.component.ml.classification

import com.feidee.fd.sml.algorithm.component.ml.MLParam
import org.apache.spark.ml.PipelineStage
import org.apache.spark.ml.classification.DecisionTreeClassifier

/**
  * @Author JunxinWang, songhaicheng
  * @Date 2019/3/21 19:07
  * @Description
  * @Reviewer YongChen
  */
case class DecisionTreeParam(
                              override val input_pt: String,
                              override val output_pt: String,
                              override val hive_table: String,
                              override val flow_time: String,
                              override val modelPath: String,
                              override var labelCol: String,
                              override val featuresCol: String,
                              override var predictionCol: String,
                              override val metrics: Array[String],
                              // 获取checkpoint的轮训次数。阈值为：>=1。当该值为-1时，则停止使用checkpoint功能,默认值为 10
                              checkpointInterval: Int,
                              // 信息增益计算方式。可选择[entropy, gini] 。默认值为："gini"
                              impurity: String,
                              // 用于离散化连续特征和选择如何在每个节点上拆分特征的最大容器数。阈值为：>= 2 。 默认值为 32
                              maxBins : Int,
                              // 树的深度。阈值：>= 0 。若深度为n，则叶节点数为 2^(n+1)-1。 默认值为：5
                              maxDepth: Int,
                              // 在树节点上考虑的分割的最小信息增益。阈值：>=0.0。 默认值为 0.0
                              minInfoGain: Double,
                              // 拆分后每个子级必须具有的最小实例数。阈值：>=1。 默认值为：1
                              minInstancesPerNode: Int,
                              // 随机种子的参数。默认值为：10
                              seed: Long,
                              // 用于预测类条件概率的列名的参数。默认值为："probability"
                              probabilityCol: String,
                              // 原始预测（a.k.a.置信度）列名称的参数。默认值："rawPrediction"
                              rawPredictionCol: String,
                              // 多类分类中阈值的参数，用于调整预测每个类的概率,数组长度和分类类别数相同
                              thresholds: Array[Double],
                              // 该算法是否缓存每个实例的节点 ID。默认值为：false
                              cacheNodeIds : Boolean
                            ) extends MLParam {

  def this() = this(null, null, null, null , null, "label", "features", "prediction", new Array[String](0),10, "gini",
    32, 5, 0.0, 1, 10, "probability", "rawPrediction", new Array[Double](0),false)

  override def verify(): Unit = {
    super.verify()
    require(Array("gini", "entropy").contains(impurity.toLowerCase),
      s"param impurity only accepts[gini, entropy], but has $impurity")
    require(maxDepth >= 0, "param maxDepth can't be negative")
    require(maxBins >= 2, "param maxBins must be not less than 2")
    require(minInstancesPerNode >= 1, "param minInstancesPerNode must be not less than 1")
    require(checkpointInterval == -1 || checkpointInterval >= 1, "param checkpointInterval must be" +
      " equals to -1 or not less than 1")
    require(minInfoGain >= 0, "param minInfoGain can't be negative")
    require(thresholds.forall(_ >= 0) && thresholds.count(_ == 0) <= 1, "Array must have length equal to the number of classes, with values > 0 excepting that at most one value may be 0")
  }

  override def toMap: Map[String, Any] = {
    var map = super.toMap
    map += ("checkpointInterval" -> checkpointInterval)
    map += ("impurity" -> impurity)
    map += ("maxBins" -> maxBins)
    map += ("maxDepth" -> maxDepth)
    map += ("minInfoGain" -> minInfoGain)
    map += ("minInstancesPerNode" -> minInstancesPerNode)
    map += ("seed" -> seed)
    map += ("probabilityCol" -> probabilityCol)
    map += ("rawPredictionCol" -> rawPredictionCol)
    map += ("thresholds" -> thresholds)
    map
  }

}


class DecisionTreeComponent extends AbstractClassificationComponent[DecisionTreeParam] {

  override def setUp(param: DecisionTreeParam): PipelineStage = {
    val dt = new DecisionTreeClassifier()
      .setCheckpointInterval(param.checkpointInterval)
      .setImpurity(param.impurity)
      .setMaxBins(param.maxBins)
      .setMaxDepth(param.maxDepth)
      .setMinInfoGain(param.minInfoGain)
      .setMinInstancesPerNode(param.minInstancesPerNode)
      .setSeed(param.seed)
      .setProbabilityCol(param.probabilityCol)
      .setRawPredictionCol(param.rawPredictionCol)
      .setThresholds(param.thresholds)
      .setFeaturesCol(param.featuresCol)
      .setPredictionCol(param.predictionCol)
      .setLabelCol(param.labelCol)
      .setCacheNodeIds(param.cacheNodeIds)

    dt
  }

}

object DecisionTreeComponent {
  def apply(paramStr: String): Unit = {
    new DecisionTreeComponent()(paramStr)
  }

  def main(args: Array[String]): Unit = {
    DecisionTreeComponent(args(0))
  }

}



